from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
from sklearn.tree import DecisionTreeClassifier
from sklearn.naive_bayes import GaussianNB
from sklearn.svm import SVC
import matplotlib.pyplot as plt
import numpy as np

iris = load_iris()
features, labels = iris.data, iris.target
classifiers = {
    'KNN': KNeighborsClassifier(),
    'Decision Tree': DecisionTreeClassifier(),
    'Naive Bayes': GaussianNB(),
    'SVM': SVC()
}
accuracies = {}

if __name__ == '__main__':
    features_train, features_test, labels_train, labels_test = train_test_split(features, labels, test_size=0.3,
                                                                                random_state=52)
    for name, clf in classifiers.items():
        clf.fit(features_train, labels_train)
        accuracy = clf.score(features_test, labels_test)
        accuracies[name] = accuracy
        print(f"算法: {name}, 分类精度: {accuracy}")

    # 选取测试点
    test_points = np.array([[5.1, 3.5, 1.4, 0.2], [6.2, 2.9, 4.3, 1.3], [7.9, 3.8, 6.4, 2.0]])
    # 绘制分类图
    plt.figure(figsize=(10, 8))
    for name, clf in classifiers.items():
        clf.fit(features[:, :2], labels)
        h = .02
        x_min, x_max = features[:, 0].min() - 1, features[:, 0].max() + 1
        y_min, y_max = features[:, 1].min() - 1, features[:, 1].max() + 1
        xx, yy = np.meshgrid(np.arange(x_min, x_max, h),
                             np.arange(y_min, y_max, h))
        Z = clf.predict(np.c_[xx.ravel(), yy.ravel()])
        Z = Z.reshape(xx.shape)
        plt.subplot(2, 2, list(classifiers.keys()).index(name) + 1)
        plt.contourf(xx, yy, Z, cmap=plt.cm.coolwarm, alpha=0.8)
        plt.scatter(features[:, 0], features[:, 1], c=labels, cmap=plt.cm.coolwarm, s=20, edgecolors='k')
        plt.scatter(test_points[:, 0], test_points[:, 1], marker='x', s=100, c='black', label='Test Points')
        plt.xlabel('Sepal length')
        plt.ylabel('Sepal width')
        plt.title(name)
    plt.tight_layout()
    plt.show()
